# -*- coding: utf-8 -*-
import os
import sys
import json
import argparse
from pathlib import Path
from typing import List

from alibabacloud_ocr_api20210707.client import Client as ocr_api20210707Client
from alibabacloud_tea_openapi import models as open_api_models
from alibabacloud_darabonba_stream.client import Client as StreamClient
from alibabacloud_ocr_api20210707 import models as ocr_api_20210707_models
from alibabacloud_tea_util import models as util_models
from alibabacloud_tea_util.client import Client as UtilClient

class OCRProcessor:
    @staticmethod
    def create_client() -> ocr_api20210707Client:
        """创建阿里云OCR客户端"""
        config = open_api_models.Config(
            access_key_id="LTAI5t6kbnkcSWY855CXbD12",
            access_key_secret="jJgVbiLFiZhpLsCBQGMF2LVxmxgfR7"
        )
        config.endpoint = 'ocr-api.cn-hangzhou.aliyuncs.com'
        return ocr_api20210707Client(config)

    @staticmethod
    def process_image(image_path: str, output_json: str) -> None:
        """处理图片并保存原始OCR结果"""
        client = OCRProcessor.create_client()
        body_stream = StreamClient.read_from_file_path(image_path)
        request = ocr_api_20210707_models.RecognizeEduPaperCutRequest(
            body=body_stream,
            cut_type='question',
            image_type='photo'
        )
        runtime = util_models.RuntimeOptions()
        try:
            resp = client.recognize_edu_paper_cut_with_options(request, runtime)
            # 将响应结果保存为JSON文件
            with open(output_json, 'w') as f:
                json.dump(resp.body.to_map(), f)
        except Exception as error:
            print(f"OCR处理失败：{error.message}")
            sys.exit(1)

class JSONProcessor:
    @staticmethod
    def remove_prefix(filename: str) -> None:
        """删除文件前6个字符"""
        with open(filename, 'r') as f:
            content = f.read()
        with open(filename, 'w') as f:
            f.write(content[6:])

    @staticmethod
    def format_json(input_path: str, output_path: str) -> None:
        """格式化JSON文件"""
        JSONProcessor.remove_prefix(input_path)
        try:
            with open(input_path, 'r', encoding='utf-8') as f:
                data = json.load(f)
            with open(output_path, 'w', encoding='utf-8') as f:
                json.dump(data, f, ensure_ascii=False, indent=4)
        except Exception as e:
            print(f"JSON格式化失败：{str(e)}")
            sys.exit(1)

class DataParser:
    @staticmethod
    def parse_and_save(input_path: str, output_path: str) -> None:
        """解析并生成最终JSON"""
        try:
            with open(input_path, 'r', encoding='utf-8') as f:
                json_data = json.load(f)
            
            data_str = json_data["body"]["Data"]
            inner_data = json.loads(data_str)
            
            formatted_data = {
                "algo_version": inner_data["algo_version"],
                "pages": [{
                    "angle": p["angle"],
                    "doc_index": p["doc_index"],
                    "height": p["height"],
                    "orgHeight": p["orgHeight"],
                    "orgWidth": p["orgWidth"],
                    "page_id": p["page_id"],
                    "subjects": [{
                        "content_list_info": s["content_list_info"],
                        "ids": s["ids"],
                        "is_multipage": s["is_multipage"],
                        "prism_wordsInfo": s["prism_wordsInfo"],
                        "text": s["text"]
                    } for s in p["subject_list"]]
                } for p in inner_data["page_list"]]
            }

            with open(output_path, 'w', encoding='utf-8') as f:
                json.dump(formatted_data, f, indent=4, ensure_ascii=False)
        except Exception as e:
            print(f"数据解析失败：{str(e)}")
            sys.exit(1)

def main():
    parser = argparse.ArgumentParser(description='试卷题目识别处理工具')
    parser.add_argument('-i', '--input', required=True, help='输入图片路径')
    parser.add_argument('-o', '--output', required=True, help='输出JSON路径')
    args = parser.parse_args()

    # 定义中间文件路径
    temp_files = {
        "raw_ocr": "Cut_output_2.json",
        "formatted": "fixed_tmp.json"
    }

    try:
        # Step 1: OCR处理
        OCRProcessor.process_image(args.input, temp_files["raw_ocr"])
        
        # Step 2: JSON预处理
        JSONProcessor.format_json(temp_files["raw_ocr"], temp_files["formatted"])
        
        # Step 3: 数据解析
        DataParser.parse_and_save(temp_files["formatted"], args.output)
        
        # 清理中间文件
        for f in temp_files.values():
            os.remove(f)
        
        print(f"✅ 处理完成，结果已保存至：{Path(args.output).resolve()}")

    except Exception as e:
        print(f"❌ 处理过程中发生错误：{str(e)}")
        sys.exit(1)

if __name__ == '__main__':
    main()
